import re
import sqlite3

import numpy as np
import pandas as pd
import streamlit as st
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaEmbeddings
from langchain_openai import OpenAI


class RetrievalLLM:
    """一个类，用于使用大语言模型（LLM）检索和重新排序offer

    参数：
        data_path (str): 包含数据csv文件的目录路径
        tables (list[str]): 数据csv文件的名称列表
        db_name (str): 用于存储数据的SQLite数据库名称

    属性：
        data_path (str): 包含数据csv文件的目录路径
        tables (list[str]): 数据csv文件的名称列表
        db_name (str): 用于存储数据的SQLite数据库名称
        llm (OpenAI): OpenAI LLM客户端
        embeddings: 嵌入模型
        db_chain (SQLDatabaseChain): 与LLM集成的SQL数据库链
        db (SQLDatabase): SQL数据库连接对象
    """

    def __init__(self, data_path, tables, db_name):
        # 初始化类属性
        self.data_path = data_path
        self.tables = tables
        self.db_name = db_name

        # 读取csv文件并存储到数据帧字典中
        dfs = {}
        for table in self.tables:
            # 从指定路径读取CSV文件
            dfs[table] = pd.read_csv(f"{self.data_path}/{table}.csv")

        # 将数据帧写入SQLite数据库
        with sqlite3.connect(self.db_name) as local_db:
            for table, df in dfs.items():
                # 将数据帧写入数据库表，如果表已存在则替换
                df.to_sql(table, local_db, if_exists='replace')

        # 创建SQL数据库连接对象，用于后续的数据库操作
        self.db = SQLDatabase.from_uri(f"sqlite:///{self.db_name}")

        # 创建LLM模型实例
        # model_name: 模型名称，留空使用默认模型
        # temperature: 控制生成文本的随机性，0表示确定性最高
        # verbose: 是否显示详细运行信息
        self.llm = OpenAI(model_name="", temperature=0, verbose=True)

        # 创建嵌入模型实例，用于文本向量化
        # base_url: Ollama服务的地址
        # model: 使用的嵌入模型名称
        self.embeddings = OllamaEmbeddings(base_url="http://127.0.0.1:11434", model="nomic-embed-text:latest")

        # 创建SQL数据库查询链，将LLM与数据库连接
        self.db_chain = create_sql_query_chain(self.llm, self.db)

    def retrieve_offers(self, prompt):
        """使用LLM从数据库中检索offer

        参数：
            prompt (str): 用于检索offer的提示词

        返回：
            list[str] or None: 检索到的offer列表，如果没有找到则返回None
        """

        # 运行SQL数据库查询链来检索offer
        retrieved_offers = self.db_chain.run(prompt)
        # 如果retrieved_offers是'None'字符串，返回None，否则返回检索到的offer
        return None if retrieved_offers == "None" else retrieved_offers

    def get_embeddings(self, documents):
        """使用嵌入模型获取文档的嵌入向量

        参数：
            documents (list[str]): 文档列表

        返回：
            np.ndarray: 包含文档嵌入向量的NumPy数组
        """

        # 如果文档列表只有一个文档，将单个文档的嵌入转换为NumPy数组
        if len(documents) == 1:
            return np.asarray(self.embeddings.embed_query(documents[0]))
        else:
            # 否则获取每个文档的嵌入并存储到列表中
            embeddings_list = []
            for document in documents:
                # 获取单个文档的嵌入向量
                embeddings_list.append(self.embeddings.embed_query(document))
            return np.asarray(embeddings_list)

    def parse_output(self, retrieved_offers, query):
        """解析retrieve_offers()方法的输出并返回一个数据帧

        参数：
            retrieved_offers (str): 检索到的offer字符串，用'#'分隔
            query (str): 用于检索offer的查询语句

        返回：
            pd.DataFrame: 包含匹配相似度和offer的数据帧，按相似度降序排列
        """

        # 使用'#'分隔符分割检索到的offer字符串
        top_offers = retrieved_offers.split("#")

        # 获取查询语句的嵌入向量
        query_embedding = self.get_embeddings([query])
        # 获取所有offer的嵌入向量
        offer_embeddings = self.get_embeddings(top_offers)

        # 计算相似度分数：
        # offer_embeddings 是一个二维的NumPy数组，包含多个offer的嵌入向量
        # query_embedding 是一个二维的NumPy数组，包含查询的嵌入向量
        # query_embedding.T 是查询嵌入的转置，使其成为一个列向量，便于矩阵乘法
        # np.dot() 计算每个offer嵌入向量与查询嵌入向量之间的点积（内积）
        # .flatten() 将二维数组转换为一维数组，得到每个offer与查询之间的相似度分数列表
        sim_scores = np.dot(offer_embeddings, query_embedding.T).flatten()
        # 将相似度得分转换为百分比形式
        sim_scores = [p * 100 for p in sim_scores]

        # 创建数据帧并按相似度排序
        df = (
            pd.DataFrame({"匹配相似度 %": sim_scores, "offer": top_offers})
            .sort_values(by=["匹配相似度 %"], ascending=False)  # 按相似度降序排列
            .reset_index(drop=True)  # 重置索引
        )
        df.index += 1  # 索引从1开始而不是0
        return df


# 数据文件路径
DATA_PATH = "data"
# 数据表名称
TABLES = ('brand_category', 'categories', 'offer_retailer')
# 数据库名称
DB_NAME = "offer_db.sqlite"
# 提示模板
PROMPT_TEMPLATE = """
你会接收到一个查询，你的任务是从`offer_retailer`表中的`OFFER`字段检索相关offer。
查询可能是混合大小写的，所以也要搜索大写版本的查询。
重要的是，你可能需要使用数据库中其他表的信息，即：`brand_category`,`categories`,`offer_retailer`,来检索正确的offer。
不要虚构offer。如果在`offer_retailer`表中找不到offer,返回字符串：`None`。
如果你能从`offer_retailer`表中检索到offer,用分隔符`#`分隔每个offer。例如，输出应该是这样的：`offer1#offer2#offer3`。
如果SQLResult为空,返回`None`。不要生成任何offer。
这是查询：`{}`
"""

# Streamlit 应用标题
st.title("搜索offer")

# 连接sqlite3数据库
conn = sqlite3.connect('offer_db.sqlite')


# 判断是否为SQL查询的函数
def is_sql_query(query):
    """判断输入的查询是否为SQL语句

    参数：
        query (str): 用户输入的查询字符串

    返回：
        bool: 如果是SQL查询返回True，否则返回False
    """
    # 定义一个包含常见SQL关键字的列表
    sql_keywords = [
        'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'DROP', 'ALTER',
        'TRUNCATE', 'MERGE', 'CALL', 'EXPLAIN', 'DESCRIBE', 'SHOW'
    ]

    # 去掉查询字符串两端的空白并转换为大写
    query_upper = query.strip().upper()

    # 遍历SQL关键字列表
    for keyword in sql_keywords:
        # 如果查询字符串以某个SQL关键字开头，返回True
        if query_upper.startswith(keyword):
            return True

    # 定义一个正则表达式模式，用于匹配以SQL关键字开头的字符串
    sql_pattern = re.compile(
        r'^\s*(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER|TRUNCATE|MERGE|CALL|EXPLAIN|DESCRIBE|SHOW)\s+',
        re.IGNORECASE  # 忽略大小写
    )

    # 如果正则表达式匹配到查询字符串，返回True
    if sql_pattern.match(query):
        return True

    # 如果查询字符串不符合任何SQL关键字模式，返回False
    return False


# 创建一个表单用于搜索
with st.form("search_form"):
    # 输入框用于输入查询
    query = st.text_input("通过类别，品牌或零售商搜索offer。")
    # 提交按钮
    submitted = st.form_submit_button("搜索")
    # 实例化RetrievalLLM类
    retrieval_llm = RetrievalLLM(DATA_PATH, TABLES, DB_NAME)

    # 如果表单提交
    if submitted:
        # 如果输入的内容是SQL语句，则显示SQL执行结果
        if is_sql_query(query):
            st.write(pd.read_sql_query(query, conn))
        # 否则使用LLM从数据库中检索offer
        else:
            # 使用提示模板格式化查询
            formatted_prompt = PROMPT_TEMPLATE.format(query)
            # 检索相关的offer
            retrieved_offers = retrieval_llm.retrieve_offers(formatted_prompt)
            # 如果没有找到相关offer
            if not retrieved_offers:
                st.text("未找到相关offer。")
            else:
                # 解析并显示检索结果
                st.table(retrieval_llm.parse_output(retrieved_offers, query))